import datetime
import zlib

from pydantic import BaseModel, Field, validator
from typing import List
from rest_framework_mongoengine.serializers import DocumentSerializer, serializers

from models.kb import KbVuln
from utilities.redis import captcha_redis as redis


class VulListSchema(BaseModel):
    """漏洞筛选条件"""

    severity: int = Field(None, title='威胁等级')
    vuln_type: str = Field(None, title='漏洞类型')
    start_time: datetime.datetime = Field(None, title='开始日期')
    end_time: datetime.datetime = Field(None, title='结束日期')
    keyword: str = Field(None, title='聚合搜索框 漏洞名称/标签/编号')

    page: int = Field(1, title='页码')
    count: int = Field(10, title='页码')


class PieVulSchema(BaseModel):
    """页面 chart 展示信息"""

    fields: List[str] = Field(..., title='需要统计的字段', min_items=1)

    @validator('fields', pre=True)
    def validate_fields(cls, fieldstr):
        values = fieldstr.split(',')
        return [v for v in values if v in ('severity', 'source')]


class LineSchema(BaseModel):
    """线性图表"""

    range: str = Field('week', title='需要统计的范围', description='week 近一周; month 近一月')

    @validator('range')
    def validate_range(cls, value):
        assert value in ('week', 'month')
        if value == 'week':
            return datetime.datetime.now() - datetime.timedelta(days=7)
        return datetime.datetime.now() - datetime.timedelta(days=30)


class PluginsListSerializer(DocumentSerializer):
    class Meta:
        model = KbVuln
        fields = ['key', 'name', 'vuln_type', 'severity']

    key = serializers.CharField(read_only=True)
    name = serializers.DictField(read_only=True)
    vuln_type = serializers.ListField(read_only=True)
    severity = serializers.IntegerField(read_only=True)


plugin_filter = {
    'is_plugin': True,
    'is_deleted': False
}


def get_plugin_dict() -> dict:
    cache_name = 'plugin_dict_cache'
    plugin_dict_cache = redis.get(cache_name)
    if not plugin_dict_cache:
        plugin_dict_cache = {}
        for item in KbVuln.objects.filter(**plugin_filter):
            value = item.to_dict()

            key = value.get('key')
            value.pop('_id', None)
            plugin_dict_cache[key] = value

        redis.set(cache_name, zlib.compress(str(plugin_dict_cache).encode()), 600)
    else:
        plugin_dict_cache = eval(zlib.decompress(plugin_dict_cache))

    return plugin_dict_cache
